/*
 * 3D Helmholtz Finite Difference Method Solver (Single Precision)
 * Solves: (I - κ∇²)u = f in 3D domain
 * Where I is identity matrix, κ is Helmholtz parameter, ∇² is Laplacian
 * 
 * The matrix is constructed to be diagonally dominant with diagonal elements near 1
 * Uses 7-point stencil for 3D finite differences
 * 
 * Uses single precision (32-bit floating point) for comparison with half precision
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <mpi.h>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

typedef struct {
    int nx, ny, nz;      // Grid dimensions
    float dx, dy, dz;    // Grid spacing
    float kappa;         // Helmholtz parameter
    float *u;            // Solution vector
    float *f;            // Right-hand side
    float *matrix;       // Coefficient matrix (7-diagonal storage)
    float *work;         // Working arrays
} HelmholtzFDM3D_Single;

// Matrix storage indices for 7-point stencil
enum {
    CENTER = 0,    // (i,j,k)
    WEST = 1,      // (i-1,j,k)  
    EAST = 2,      // (i+1,j,k)
    SOUTH = 3,     // (i,j-1,k)
    NORTH = 4,     // (i,j+1,k)
    BOTTOM = 5,    // (i,j,k-1)
    TOP = 6        // (i,j,k+1)
};

// Convert 3D indices to linear index
static inline int idx3d(int i, int j, int k, int nx, int ny, int nz) {
    return i + nx * (j + ny * k);
}

// Initialize the Helmholtz FDM solver
HelmholtzFDM3D_Single* helmholtz_fdm_single_init(int nx, int ny, int nz, 
                                                  float lx, float ly, float lz,
                                                  float kappa) {
    HelmholtzFDM3D_Single* solver = malloc(sizeof(HelmholtzFDM3D_Single));
    
    solver->nx = nx;
    solver->ny = ny; 
    solver->nz = nz;
    solver->dx = lx / (float)(nx - 1);
    solver->dy = ly / (float)(ny - 1);
    solver->dz = lz / (float)(nz - 1);
    solver->kappa = kappa;
    
    int total_points = nx * ny * nz;
    
    // Allocate memory
    solver->u = calloc(total_points, sizeof(float));
    solver->f = calloc(total_points, sizeof(float));
    solver->matrix = calloc(7 * total_points, sizeof(float));
    solver->work = calloc(total_points, sizeof(float));
    
    if (!solver->u || !solver->f || !solver->matrix || !solver->work) {
        fprintf(stderr, "Memory allocation failed\n");
        exit(1);
    }
    
    return solver;
}

// Set up the coefficient matrix for (I - κ∇²)u = f
void helmholtz_fdm_single_setup_matrix(HelmholtzFDM3D_Single* solver) {
    int nx = solver->nx, ny = solver->ny, nz = solver->nz;
    float dx = solver->dx, dy = solver->dy, dz = solver->dz;
    float kappa = solver->kappa;
    
    // Finite difference coefficients for second derivatives
    float cx = kappa / (dx * dx);
    float cy = kappa / (dy * dy);
    float cz = kappa / (dz * dz);
    
    printf("Setting up matrix with kappa=%.6f, dx=%.6f, dy=%.6f, dz=%.6f\n", 
           kappa, dx, dy, dz);
    printf("Coefficients: cx=%.6f, cy=%.6f, cz=%.6f\n", cx, cy, cz);
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                
                // Initialize all stencil points to zero
                for (int m = 0; m < 7; m++) {
                    solver->matrix[7 * idx + m] = 0.0f;
                }
                
                // Interior points use full 7-point stencil
                if (i > 0 && i < nx-1 && j > 0 && j < ny-1 && k > 0 && k < nz-1) {
                    // Center point: 1 + κ(2/dx² + 2/dy² + 2/dz²) for (I - κ∇²)
                    solver->matrix[7 * idx + CENTER] = 1.0f + 2.0f * (cx + cy + cz);
                    
                    // Neighbor points: -κ/h²
                    solver->matrix[7 * idx + WEST]   = -cx;
                    solver->matrix[7 * idx + EAST]   = -cx;
                    solver->matrix[7 * idx + SOUTH]  = -cy;
                    solver->matrix[7 * idx + NORTH]  = -cy;
                    solver->matrix[7 * idx + BOTTOM] = -cz;
                    solver->matrix[7 * idx + TOP]    = -cz;
                }
                // Boundary points: Homogeneous Dirichlet (u = 0)
                else {
                    solver->matrix[7 * idx + CENTER] = 1.0f;
                    // All other coefficients remain zero for boundary
                }
            }
        }
    }
    
    printf("Matrix setup complete. Sample diagonal element: %.6f\n", 
           solver->matrix[7 * idx3d(nx/2, ny/2, nz/2, nx, ny, nz) + CENTER]);
}

// Set right-hand side with analytical test function
void helmholtz_fdm_single_set_rhs(HelmholtzFDM3D_Single* solver) {
    int nx = solver->nx, ny = solver->ny, nz = solver->nz;
    float dx = solver->dx, dy = solver->dy, dz = solver->dz;
    float kappa = solver->kappa;
    
    // Test function: u(x,y,z) = sin(πx/Lx) * sin(πy/Ly) * sin(πz/Lz)
    // For this function, -∇²u = π²(1/Lx² + 1/Ly² + 1/Lz²) * u
    float lx = (float)(nx - 1) * dx;
    float ly = (float)(ny - 1) * dy;
    float lz = (float)(nz - 1) * dz;
    
    float lambda = (float)(M_PI * M_PI) * (1.0f/(lx*lx) + 1.0f/(ly*ly) + 1.0f/(lz*lz));
    
    printf("Setting RHS with lambda=%.6f, domain size: %.3f x %.3f x %.3f\n", 
           lambda, lx, ly, lz);
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                
                float x = (float)i * dx;
                float y = (float)j * dy; 
                float z = (float)k * dz;
                
                // u(x,y,z) = sin(πx/Lx) * sin(πy/Ly) * sin(πz/Lz)
                float u_exact = sinf((float)M_PI * x / lx) * sinf((float)M_PI * y / ly) * sinf((float)M_PI * z / lz);
                
                // (I - κ∇²)u = u + κλu = (1 + κλ)u
                if (i == 0 || i == nx-1 || j == 0 || j == ny-1 || k == 0 || k == nz-1) {
                    // Boundary: u = 0
                    solver->f[idx] = 0.0f;
                } else {
                    // Interior: f = (1 + κλ)u
                    solver->f[idx] = (1.0f + kappa * lambda) * u_exact;
                }
            }
        }
    }
    
    printf("RHS setup complete. Max RHS value: %.6f\n", 
           solver->f[idx3d(nx/2, ny/2, nz/2, nx, ny, nz)]);
}

// Matrix-vector multiplication: y = A*x
void helmholtz_fdm_single_matvec(HelmholtzFDM3D_Single* solver, float* x, float* y) {
    int nx = solver->nx, ny = solver->ny, nz = solver->nz;
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                float* coeff = &solver->matrix[7 * idx];
                
                y[idx] = coeff[CENTER] * x[idx];
                
                // Add neighbor contributions if they exist
                if (i > 0)    y[idx] += coeff[WEST]   * x[idx3d(i-1, j, k, nx, ny, nz)];
                if (i < nx-1) y[idx] += coeff[EAST]   * x[idx3d(i+1, j, k, nx, ny, nz)];
                if (j > 0)    y[idx] += coeff[SOUTH]  * x[idx3d(i, j-1, k, nx, ny, nz)];
                if (j < ny-1) y[idx] += coeff[NORTH]  * x[idx3d(i, j+1, k, nx, ny, nz)];
                if (k > 0)    y[idx] += coeff[BOTTOM] * x[idx3d(i, j, k-1, nx, ny, nz)];
                if (k < nz-1) y[idx] += coeff[TOP]    * x[idx3d(i, j, k+1, nx, ny, nz)];
            }
        }
    }
}

// Compute vector 2-norm
float vector_norm_single(float* v, int n) {
    float sum = 0.0f;
    for (int i = 0; i < n; i++) {
        sum += v[i] * v[i];
    }
    return sqrtf(sum);
}

// Simple Jacobi iterative solver
int helmholtz_fdm_single_solve_jacobi(HelmholtzFDM3D_Single* solver, float tol, int max_iter) {
    int n = solver->nx * solver->ny * solver->nz;
    float* u_old = solver->work;
    
    printf("Starting Jacobi solver with tol=%.2e, max_iter=%d\n", tol, max_iter);
    
    // Initialize solution to zero
    memset(solver->u, 0, n * sizeof(float));
    
    float final_residual = 1e10f;
    
    for (int iter = 0; iter < max_iter; iter++) {
        // Copy current solution
        memcpy(u_old, solver->u, n * sizeof(float));
        
        // Jacobi update: u_new[i] = (f[i] - A_off[i]*u_old) / A_diag[i]
        for (int k = 0; k < solver->nz; k++) {
            for (int j = 0; j < solver->ny; j++) {
                for (int i = 0; i < solver->nx; i++) {
                    int idx = idx3d(i, j, k, solver->nx, solver->ny, solver->nz);
                    float* coeff = &solver->matrix[7 * idx];
                    
                    float rhs = solver->f[idx];
                    
                    // Subtract off-diagonal terms
                    if (i > 0)           rhs -= coeff[WEST]   * u_old[idx3d(i-1, j, k, solver->nx, solver->ny, solver->nz)];
                    if (i < solver->nx-1) rhs -= coeff[EAST]   * u_old[idx3d(i+1, j, k, solver->nx, solver->ny, solver->nz)];
                    if (j > 0)           rhs -= coeff[SOUTH]  * u_old[idx3d(i, j-1, k, solver->nx, solver->ny, solver->nz)];
                    if (j < solver->ny-1) rhs -= coeff[NORTH]  * u_old[idx3d(i, j+1, k, solver->nx, solver->ny, solver->nz)];
                    if (k > 0)           rhs -= coeff[BOTTOM] * u_old[idx3d(i, j, k-1, solver->nx, solver->ny, solver->nz)];
                    if (k < solver->nz-1) rhs -= coeff[TOP]    * u_old[idx3d(i, j, k+1, solver->nx, solver->ny, solver->nz)];
                    
                    // Update solution
                    solver->u[idx] = rhs / coeff[CENTER];
                }
            }
        }
        
        // Check convergence every 10 iterations
        if (iter % 10 == 0) {
            // Compute residual: r = f - A*u
            helmholtz_fdm_single_matvec(solver, solver->u, solver->work);
            for (int i = 0; i < n; i++) {
                solver->work[i] = solver->f[i] - solver->work[i];
            }
            
            float residual_norm = vector_norm_single(solver->work, n);
            final_residual = residual_norm;
            printf("Iteration %d: residual norm = %.6e\n", iter, residual_norm);
            
            if (residual_norm < tol) {
                printf("Converged in %d iterations\n", iter);
                return iter;
            }
        }
    }
    
    printf("Maximum iterations reached. Final residual: %.6e\n", final_residual);
    return max_iter;
}

// Compute analytical solution for verification
void helmholtz_fdm_single_analytical_solution(HelmholtzFDM3D_Single* solver, float* u_exact) {
    int nx = solver->nx, ny = solver->ny, nz = solver->nz;
    float dx = solver->dx, dy = solver->dy, dz = solver->dz;
    
    float lx = (float)(nx - 1) * dx;
    float ly = (float)(ny - 1) * dy;
    float lz = (float)(nz - 1) * dz;
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                
                float x = (float)i * dx;
                float y = (float)j * dy;
                float z = (float)k * dz;
                
                // Boundary points have u = 0
                if (i == 0 || i == nx-1 || j == 0 || j == ny-1 || k == 0 || k == nz-1) {
                    u_exact[idx] = 0.0f;
                } else {
                    u_exact[idx] = sinf((float)M_PI * x / lx) * sinf((float)M_PI * y / ly) * sinf((float)M_PI * z / lz);
                }
            }
        }
    }
}

// Cleanup memory
void helmholtz_fdm_single_destroy(HelmholtzFDM3D_Single* solver) {
    if (solver) {
        free(solver->u);
        free(solver->f);
        free(solver->matrix);
        free(solver->work);
        free(solver);
    }
}

// Test function
int test_helmholtz_fdm_3d_single() {
    printf("=== 3D Helmholtz FDM Solver Test (Single Precision) ===\n");
    
    // Problem parameters
    int nx = 32, ny = 32, nz = 32;
    float lx = 1.0f, ly = 1.0f, lz = 1.0f;
    float kappa = 0.01f;  // Smaller value for better diagonal dominance
    
    printf("Grid: %dx%dx%d, Domain: %.1fx%.1fx%.1f, kappa=%.3f\n", 
           nx, ny, nz, lx, ly, lz, kappa);
    
    // Initialize solver
    HelmholtzFDM3D_Single* solver = helmholtz_fdm_single_init(nx, ny, nz, lx, ly, lz, kappa);
    
    // Setup matrix and RHS
    helmholtz_fdm_single_setup_matrix(solver);
    helmholtz_fdm_single_set_rhs(solver);
    
    // Solve the system
    int iterations = helmholtz_fdm_single_solve_jacobi(solver, 1e-8f, 1000);
    
    // Verify against analytical solution
    float* u_exact = malloc(nx * ny * nz * sizeof(float));
    helmholtz_fdm_single_analytical_solution(solver, u_exact);
    
    // Compute error
    float error = 0.0f;
    for (int i = 0; i < nx * ny * nz; i++) {
        float diff = solver->u[i] - u_exact[i];
        error += diff * diff;
    }
    error = sqrtf(error) / sqrtf((float)(nx * ny * nz));
    
    printf("L2 error vs analytical solution: %.6e\n", error);
    printf("Solution norm: %.6e\n", vector_norm_single(solver->u, nx * ny * nz));
    printf("Exact solution norm: %.6e\n", vector_norm_single(u_exact, nx * ny * nz));
    printf("Iterations used: %d\n", iterations);
    
    // Check matrix diagonal dominance
    int dominant_count = 0;
    float min_diag = 1e10f, max_diag = -1e10f;
    for (int k = 1; k < nz-1; k++) {
        for (int j = 1; j < ny-1; j++) {
            for (int i = 1; i < nx-1; i++) {
                int idx = idx3d(i, j, k, nx, ny, nz);
                float* coeff = &solver->matrix[7 * idx];
                
                float diag = fabsf(coeff[CENTER]);
                float off_diag_sum = fabsf(coeff[WEST]) + fabsf(coeff[EAST]) + 
                                     fabsf(coeff[SOUTH]) + fabsf(coeff[NORTH]) + 
                                     fabsf(coeff[BOTTOM]) + fabsf(coeff[TOP]);
                
                if (diag > off_diag_sum) dominant_count++;
                
                min_diag = fminf(min_diag, coeff[CENTER]);
                max_diag = fmaxf(max_diag, coeff[CENTER]);
            }
        }
    }
    
    int total_interior = (nx-2) * (ny-2) * (nz-2);
    printf("Diagonally dominant points: %d/%d (%.1f%%)\n", 
           dominant_count, total_interior, 100.0 * dominant_count / total_interior);
    printf("Diagonal element range: [%.6f, %.6f]\n", min_diag, max_diag);
    
    // Success criteria: good error
    int success = (error < 1e-4f);  // Reasonable tolerance for single precision
    printf("Test result: %s (error=%.2e < 1e-4)\n", success ? "PASSED" : "FAILED", error);
    
    free(u_exact);
    helmholtz_fdm_single_destroy(solver);
    
    return success ? 0 : 1;  // Return 0 for success
}

// Main function for standalone testing
#ifdef HELMHOLTZ_FDM_SINGLE_MAIN
int main(int argc, char** argv) {
    return test_helmholtz_fdm_3d_single();
}
#endif 